import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans


def load_file(file_path):
    """
    从文件中加载数据

    :param file_path: 文件路径
    :return: 文件数据 (numpy.ndarray) 或 None
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            record = file.read()
        record_list = [list(map(float, line.split('\t'))) for line in record.splitlines() if line.strip()]
        return np.array(record_list)
    except FileNotFoundError:
        print("文件不存在, 请检查路径是否正确: ", file_path)
        return None
    except Exception as e:
        print("发生错误: ", e)
        return None


def euclidean_distance(vector_a, vector_b):
    """
    计算两向量之间的欧氏距离

    :param vector_a: (numpy.ndarray) 向量 A
    :param vector_b: (numpy.ndarray) 向量 B
    :return: 向量 A 与向量 B 之间的欧式距离
    """
    return np.linalg.norm(vector_a - vector_b, ord=2)


def initialize_centroids(dataset, num_centroids):
    """
    初始化聚类中心

    :param dataset: 数据集 (numpy.ndarray), 每行代表一个样本, 每列代表一个特征
    :param num_centroids: 聚类中心的数量
    :return: 聚类中心数组, 每行代表一个聚类中心, 每列代表一个特征
    """
    num_features = np.shape(dataset)[1]
    centroids = np.zeros((num_centroids, num_features))
    for feature_index in range(num_features):
        # 获取当前特征维度上的最小值和最大值
        min_value = min(dataset[:, feature_index])
        max_value = max(dataset[:, feature_index])
        # 利用随机函数生成一个形状为 (k, 1) 的随机数组, 取值范围在最小值和最大值之间
        centroids[:, feature_index] = min_value + (max_value - min_value) * np.random.rand(num_centroids)
    return centroids


def kmeans(dataset, num_clusters):
    """
    kmeans 聚类

    :param dataset: 数据集 (numpy.ndarray), 每行代表一个样本, 每列代表一个特征
    :param num_clusters: 簇的数量
    :return: 聚类中心 (numpy.ndarray), 聚类分配信息 (numpy.ndarray)
    """
    num_samples, num_features = dataset.shape  # 获取数据集中的样本数量和特征维度
    centroids = initialize_centroids(dataset, num_clusters)
    cluster_info = np.zeros((num_samples, 2))
    cluster_changed = True
    while cluster_changed:
        cluster_changed = False
        for sample_index in range(num_samples):
            # 计算当前样本到每个聚类中心的距离, 并找到最近的聚类中心
            distances_to_centroids = [euclidean_distance(dataset[sample_index, :], centroids[cluster_index, :]) for
                                      cluster_index in range(num_clusters)]
            min_distance = min(distances_to_centroids)
            min_cluster_index = distances_to_centroids.index(min_distance)
            # 若当前样本的聚类中心发生了变化, 则对标志变量进行更新
            if cluster_info[sample_index, 0] != min_cluster_index:
                cluster_changed = True
            cluster_info[sample_index, :] = min_cluster_index, min_distance
        # 更新聚类中心
        for cluster_index in range(num_clusters):
            samples_in_cluster = dataset[np.nonzero(cluster_info[:, 0] == cluster_index)[0]]
            centroids[cluster_index, :] = np.mean(samples_in_cluster, axis=0)
    return centroids, cluster_info


def plot_clusters(dataset, centroids, cluster_info):
    """
    绘制聚类结果的散点图

    :param dataset: 数据集 (numpy.ndarray), 每行代表一个样本, 每列代表一个特征
    :param centroids: 聚类中心 (numpy.ndarray)
    :param cluster_info: 聚类分配信息 (numpy.ndarray)
    """
    num_clusters = centroids.shape[0]
    colors = ['r', 'g', 'b', 'c', 'm', 'k']
    # 绘制样本点
    for sample_index in range(dataset.shape[0]):
        cluster_index = int(cluster_info[sample_index, 0])
        plt.scatter(dataset[sample_index, 0], dataset[sample_index, 1], c=colors[cluster_index], alpha=0.5)
    # 绘制聚类中心
    for cluster_index in range(num_clusters):
        plt.scatter(centroids[cluster_index, 0], centroids[cluster_index, 1], c='y', s=100)

    plt.xlabel('Feature 1')
    plt.ylabel('Feature 2')
    plt.title('Manually Implemented KMeans Clustering Result')
    plt.grid(True)
    plt.show()


if __name__ == '__main__':
    dataset = load_file('TESTDATA.TXT')
    # 手动实现 KMeans 均值算法
    centroids, cluster_info = kmeans(dataset, 3)
    plot_clusters(dataset, centroids, cluster_info)
    # sklearn 库中的 KMeans 均值算法
    kmeans = KMeans(n_clusters=3, random_state=42)
    kmeans.fit(dataset)
    cluster_labels = kmeans.labels_
    centroids = kmeans.cluster_centers_
    plt.scatter(dataset[:, 0], dataset[:, 1], c=cluster_labels, alpha=0.5)
    plt.scatter(centroids[:, 0], centroids[:, 1], c='y', s=100)
    plt.xlabel('Feature 1')
    plt.ylabel('Feature 2')
    plt.title('Sklearn KMeans Clustering Result')
    plt.grid(True)
    plt.show()
